package zip;

import static zip.Utils.to;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.EOFException;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.util.PriorityQueue;
import java.util.Queue;

public class HuffZip {
	public static void main(String[] args) {
		doHuffZip("D:\\workspace\\Java workspace\\lanJie\\src\\zip\\tests\\test.txt");
	}

	private static Queue<Node> pq = null; // priorityQueue用于构造Huffman树时排序
	private static int[] freq = null; // 记录每个字符的出现次数
	private static String[] codes = null; // 记录Huffman编码
	private static Node root = null; // Huffman树根
	private static int bytesCount; // 用于记录一共有多少有效字符
	private static final int CHS = 300; // 总的字符种类有256个

	public static class Node implements Comparable<Node> {
		Node right, left;
		String code;
		int freq, key; // key为字符的ASCII码值

		public Node(Node l, Node r, int freq) {
			this.right = r;
			this.left = l;
			this.freq = freq;
		}

		@Override
		public int compareTo(Node o) {
			return this.freq - o.freq;
		}
	}

	/**
	 * 压缩文件对外的接口方法
	 * 
	 * @param pathName 要压缩的文件路径
	 */
	public static void doHuffZip(String pathName) {
		initZip();
		calFreq(pathName);
		createTree();
		realZip(midZip(pathName), pathName + ".huff");
	}

	private static void initZip() {
		pq = new PriorityQueue<>();
		freq = new int[CHS];
		codes = new String[CHS];
		root = null;
		bytesCount = 0;
	}

	/**
	 * 统计文本每一个byte的出现次数
	 * 
	 * @param pathName 输入文本路径名
	 */
	private static void calFreq(String pathName) {
		File file = new File(pathName);
		try {
			DataInputStream din = new DataInputStream(new FileInputStream(file));
			while (true) {
				try {
					freq[to(din.readByte())] += 1;
				} catch (EOFException eof) {
					break;
				}
			}
			din.close();
		} catch (IOException ioe) {
			System.out.println("IO Exception = " + ioe);
		}
		file = null;
	}

	/**
	 * 构建huffman树
	 */
	private static void createTree() {
		for (int i = 0; i < CHS; ++i) {
			if (freq[i] > 0) {
				Node node = new Node(null, null, freq[i]);
				node.key = i;
				pq.add(node);
				bytesCount += 1;
				// System.out.println(i + ":" + freq[i] + " ");
			}
		}
		// System.out.println("");
		// System.out.println("pq size = " + pq.size());
		if (bytesCount == 0)
			return;
		while (pq.size() > 1) {
			Node temp1 = pq.poll(), temp2 = pq.poll();
			Node fath = new Node(temp1, temp2, temp1.freq + temp2.freq);
			pq.add(fath);
		}
		root = pq.poll();
		setCodes(root, "");
		// System.out.println("total value = " + root.freq);
	}

	/**
	 * 前序遍历树，获取字符的哈夫曼编码
	 * 
	 * @param root  树根
	 * @param code  当前构建的哈夫曼编码
	 * @param codes 存储哈夫曼编码的数组
	 */
	private static void setCodes(Node root, String code) {
		if (root == null)
			return;
		root.code = code;
		if (root.left == null && root.right == null)
			codes[root.key] = code;
		setCodes(root.left, code + "0");
		setCodes(root.right, code + "1");
	}

	/**
	 * 压缩时临时产生辅助的中间文件 </br>
	 * 文件中以字符串的形式记录哈夫曼编码 所以物理上的一个byte存储逻辑上的一个bit
	 * 
	 * @param pathName 要压缩的文件路径
	 * @return 中间文件的路径名
	 */
	private static String midZip(String pathName) {
		String midpath = pathName + ".mid";
		File midFile = new File(midpath);
		try {
			DataInputStream din = new DataInputStream(new FileInputStream(new File(pathName)));
			PrintStream printStream = new PrintStream(midFile);

			while (true) {
				try {
					byte bt = din.readByte();
					printStream.print(codes[to(bt)]);
				} catch (EOFException eof) {
					break;
				}
			}
			din.close();
			printStream.close();
		} catch (IOException ioe) {
			System.out.println("IO Exception = " + ioe);
		}
		return midpath;
	}

	/**
	 * 将01字符串压缩成01位串
	 * 
	 * @param midPath 中间文件，存储的是01字符串
	 * @param tarPath 目标输出文件，存储压缩之后的结果
	 * @apiNote 目标文件的存储格式：</br>
	 *          总的有效字符数[int] + (有效字符[byte] + 频数[int])s +</br>
	 *          多余的凑数bit个数[int] + 哈夫曼编码[bits]s
	 */
	private static void realZip(String midPath, String tarPath) {
		File inFile = new File(midPath);
		File outFile = new File(tarPath);
		try {
			DataInputStream din = new DataInputStream(new FileInputStream(inFile));
			DataOutputStream dout = new DataOutputStream(new FileOutputStream(outFile));

			dout.write(bytesCount);
			// System.out.println("bytes counter = " + bytesCount);
			for (int i = 0; i < (1 << 8); ++i) {
				if (freq[i] > 0) {
					byte ch = (byte) i;
					dout.write(ch);
					dout.writeInt(freq[i]);
				}
			}
			
			byte extraBitsCount = (byte) (int) ((8L - (inFile.length() % 8L)) % 8L);
			dout.write(extraBitsCount);
			// System.out.println("extra bits count = " + extraBitsCount);

			// 将01字符串压缩成位串
			// 由于Java中最小能够直接表示的类型为byte
			// 所以利用位操作
			byte bits = 0, ch;
			int iter = 0;
			while (true) {
				try {
					for (iter = 0; iter < 8; ++iter) {
						ch = din.readByte();
						bits <<= 1;
						if (ch == '1')
							bits |= 1;
					}
					dout.write(bits);
				} catch (EOFException eof) {
					if (extraBitsCount > 0) {
						bits <<= extraBitsCount;
						dout.write(bits);
					}
					break;
				}
			}
			din.close();
			dout.close();
			// System.out.println("the Output File Size = " + outFile.length());

		} catch (IOException ioe) {
			System.out.println("IO Exception = " + ioe);
		}
		inFile.delete(); // 删除中间文件
	}
}
